import pickle
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from torch.utils.data import DataLoader, TensorDataset
import json
import matplotlib.pyplot as plt 

# 加载 annotation.json 文件
with open('data/iu_xray/iu_xray/embeddings.json', 'r') as f:
    annotation = json.load(f)

# 提取图像和报告的嵌入向量
image_embeddings = []
text_embeddings = []
for item in annotation['train']:
    image_embeddings.append(item['image_embedding'])
    text_embeddings.append(item['report_embedding'])

# 转换为 numpy 数组
image_embeddings = np.array(image_embeddings)
text_embeddings = np.array(text_embeddings)

# 定义嵌入维度
embedding_dim = image_embeddings.shape[1]

class CVAE(nn.Module):
    def __init__(self, input_dim, latent_dim):
        super(CVAE, self).__init__()
        self.latent_dim = latent_dim
        
        # 编码器部分
        self.encoder = nn.Sequential(
            nn.Linear(input_dim + input_dim, 256),
            nn.ReLU(),
            nn.Linear(256, 128),
            nn.ReLU()
        )
        self.fc_mu = nn.Linear(128, latent_dim)
        self.fc_logvar = nn.Linear(128, latent_dim)
        
        # 解码器部分
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim + input_dim, 128),
            nn.ReLU(),
            nn.Linear(128, 256),
            nn.ReLU(),
            nn.Linear(256, input_dim)
        )
    
    def encode(self, x, c):
        # 将输入和条件拼接
        x = torch.cat([x, c], dim=1)
        h = self.encoder(x)
        mu = self.fc_mu(h)
        logvar = self.fc_logvar(h)
        return mu, logvar
    
    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std
    
    def decode(self, z, c):
        # 将潜在变量和条件拼接
        z = torch.cat([z, c], dim=1)
        return self.decoder(z)
    
    def forward(self, x, c):
        mu, logvar = self.encode(x, c)
        z = self.reparameterize(mu, logvar)
        return self.decode(z, c), mu, logvar
    
def loss_function(recon_x, x, mu, logvar):
    BCE = F.mse_loss(recon_x, x, reduction='sum')
    KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return BCE + KLD
# 重构误差 KL散度


def use():
    # 转换为 Tensor
    image_embeddings_tensor = torch.tensor(image_embeddings, dtype=torch.float32)
    text_embeddings_tensor = torch.tensor(text_embeddings, dtype=torch.float32)

    # 创建数据集
    dataset = TensorDataset(image_embeddings_tensor, text_embeddings_tensor)
    train_size = int(0.95 * len(dataset))
    test_size = len(dataset) - train_size
    train_dataset, test_dataset = torch.utils.data.random_split(dataset, [train_size, test_size])

    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

    # 初始化模型、优化器
    model = CVAE(input_dim=embedding_dim, latent_dim=32)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

    # 训练模型
    num_epochs = 60
    loss_history = []
    for epoch in range(num_epochs):
        model.train()
        train_loss = 0
        for batch_idx, (image_emb, text_emb) in enumerate(train_loader):
            optimizer.zero_grad()
            recon_batch, mu, logvar = model(image_emb, text_emb)
            loss = loss_function(recon_batch, text_emb, mu, logvar)
            loss.backward()
            train_loss += loss.item()
            optimizer.step()
        
        avg_loss = train_loss / len(train_loader)
        loss_history.append(avg_loss)  # 添加这一行
        print(f'Epoch {epoch+1}/{num_epochs}, Loss: {train_loss/len(train_loader.dataset):.4f}')

    # 保存模型
    # 在训练结束后，绘制损失曲线
    plt.plot(loss_history, label='Training Loss')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.title('Training Loss Curve')
    plt.legend()
    plt.show()
    torch.save(model.state_dict(), 'src/model_pth/cvae_model_attention.pth')

if __name__ == "__main__":
    use()